# -*- coding: utf-8 -*-
"""
Created on Thu Apr 21 23:27:29 2022

@author: Jovi Wong
"""
import os
import torch
import dgl.function as dglfn
from SignNetGraph import SignNetGraph

class node_diff_to_edge():
    def __init__(self, src_kw:str, dst_kw:str, edge_kw:str, decay_rate:float=0.9):
        """
        Collect information from edge-connected nodes to calculate the displacement direction
        """
        self.src_kw = src_kw
        self.dst_kw = dst_kw
        self.edge_kw = edge_kw
        self.decay_rate = decay_rate

    def __call__(self, edge):            
        edge.data[self.edge_kw] = edge.data[self.edge_kw] * self.decay_rate
        diff = edge.src[self.src_kw] - edge.dst[self.dst_kw]
        disp = torch.zeros(diff.shape)
        for idx in range(len(edge.data[self.edge_kw])):
            if edge.data[self.edge_kw][idx] > 0:
                disp[idx] = edge.data[self.edge_kw][idx] * diff[idx]
            else:
                disp[idx] = edge.data[self.edge_kw][idx] * diff[idx]
        return {"displacement": disp}

class Augmentor():
    def __init__(self, raw_data_path, feat_path):
        dataset_dir, dataset_name = os.path.split(raw_data_path)
        graph_name = dataset_name.split(".")[0]
        self.graph_dataset = SignNetGraph(graph_name, raw_data_path, feat_path)
        self.graph = self.graph_dataset[0]

    def train(self, epochs:int):
        for epoch in range(epochs):
            print("<<<<< Augmenting {0}/{1} epochs>>>>>".format(epoch+1, epochs))
            self.graph.apply_edges(func=node_diff_to_edge("feat", "feat", "sign"))
            self.graph.update_all(message_func=dglfn.u_add_e("feat", "displacement","m"), 
                                  reduce_func=dglfn.sum("m", "feat"))
            return self.graph_dataset.graph_path